{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "17d28666-31e5-4282-a0ff-a0a69270f4e9",
   "metadata": {},
   "source": [
    "## SFTTrainer Example - Finetuning with Sagemaker\n",
    "\n",
    "This notebook demonstrates basic user flow for SFT Finetuning from a model available in Sagemaker Jumpstart.\n",
    "Information on available models on jumpstart: https://docs.aws.amazon.com/sagemaker/latest/dg/jumpstart-foundation-models-latest.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "635da80f-e82d-4434-b0ac-9d8ac4600ffb",
   "metadata": {},
   "source": [
    "### Setup and Configuration\n",
    "\n",
    "Initialize the environment by importing necessary libraries and configuring AWS credentials"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# Configure AWS credentials and region\n",
    "#! ada credentials update --provider=isengard --account=<> --role=Admin --profile=default --once\n",
    "#! aws configure set region us-west-2"
   ],
   "id": "87aa2004556ad7c6",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "a51be0b5-fd33-4fa0-af2b-d08ce0dc7a8e",
   "metadata": {},
   "source": [
    "from sagemaker.train.sft_trainer import SFTTrainer\n",
    "from sagemaker.train.common import TrainingType\n",
    "from sagemaker.core.training.configs import InputData\n",
    "from rich import print as rprint\n",
    "from rich.pretty import pprint\n",
    "from sagemaker.core.resources import ModelPackage\n",
    "\n",
    "import boto3\n",
    "from sagemaker.core.helper.session_helper import Session\n",
    "import os\n",
    "\n",
    "\n",
    "# For MLFlow native metrics in Trainer wait, run below line with approriate region\n",
    "os.environ[\"SAGEMAKER_MLFLOW_CUSTOM_ENDPOINT\"] = \"https://mlflow.sagemaker.us-west-2.app.aws\"\n",
    "\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "96994fe5-c9cf-44ab-8f52-be895ca289ab",
   "metadata": {},
   "source": [
    "## Finetuning with Jumpstart base model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18c4a2a3-2a6d-44c0-a1a7-14938bf2ff83",
   "metadata": {},
   "source": [
    "### Create SFTTrainer\n",
    "**Required Parameters** \n",
    "\n",
    "* `model`: base_model id on Sagemaker Hubcontent that is available to finetune (or) ModelPackage artifacts\n",
    "\n",
    "**Optional Parameters**\n",
    "* `training_type`: Choose from TrainingType Enum(sagemaker.modules.train.common) either LORA OR FULL.\n",
    "* `model_package_group_name`: ModelPackage group name or ModelPackageGroup\n",
    "* `mlflow_resource_arn`: MLFlow app ARN to track the training job\n",
    "* `mlflow_experiment_name`: MLFlow app experiment name(str)\n",
    "* `mlflow_run_name`: MLFlow app run name(str)\n",
    "* `training_dataset`: Training Dataset - either Dataset ARN or S3 Path of the dataset (Please note these are required for a training job to run, can be either provided via Trainer or .train())\n",
    "* `validation_dataset`: Validation Dataset - either Dataset ARN or S3 Path of the dataset\n",
    "* `s3_output_path`: S3 path for the trained model artifacts"
   ]
  },
  {
   "cell_type": "code",
   "id": "88fe8360-de50-481d-932f-564a32be66a0",
   "metadata": {},
   "source": [
    "# For fine-tuning \n",
    "sft_trainer = SFTTrainer(\n",
    "    model=\"meta-textgeneration-llama-3-2-1b-instruct\", # Union[str, ModelPackage]\n",
    "    training_type=TrainingType.LORA, \n",
    "    model_package_group_name=\"arn:aws:sagemaker:us-west-2:<>:model-package-group/sdk-test-finetuned-models\", # Make it Optional\n",
    "    #mlflow_resource_arn=\"arn:aws:sagemaker:us-west-2:<>:mlflow-tracking-server/mmlu-eval-experiment\",  # Optional[str] - MLflow app ARN (auto-resolved if not provided), can accept name and search in the account\n",
    "    mlflow_experiment_name=\"test-finetuned-models-exp\", # Optional[str]\n",
    "    mlflow_run_name=\"test-finetuned-models-run\", # Optional[str]\n",
    "    training_dataset=\"s3://mc-flows-sdk-testing/input_data/sft/\", #Optional[]\n",
    "    s3_output_path=\"s3://mc-flows-sdk-testing/output/\",\n",
    "    accept_eula=True\n",
    ")\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "9325b23b-d6a5-425e-b588-77d3b97f3843",
   "metadata": {},
   "source": [
    "# For fine-tuning (us-east-1)\n",
    "sft_trainer = SFTTrainer(\n",
    "    model=\"meta-textgeneration-llama-3-2-1b-instruct\", # Union[str, ModelPackage]\n",
    "    training_type=TrainingType.LORA, \n",
    "    model_package_group_name=\"arn:aws:sagemaker:us-east-1:<>:model-package-group/sdk-test-finetuned-models-us-east-1\", # Make it Optional\n",
    "    mlflow_resource_arn=\"arn:aws:sagemaker:us-east-1:<>:mlflow-app/app-J2NPD6IV77BJ\",  # Optional[str] - MLflow app ARN (auto-resolved if not provided), can accept name and search in the account\n",
    "    mlflow_experiment_name=\"test-finetuned-models-exp\", # Optional[str]\n",
    "    mlflow_run_name=\"test-finetuned-models-run\", # Optional[str]\n",
    "    training_dataset=\"s3://mc-flows-sdk-testing-us-east-1/input_data/sft/train_285.jsonl\", #Optional[]\n",
    "    s3_output_path=\"s3://mc-flows-sdk-testing-us-east-1/output/\",\n",
    "    #sagemaker_session=sagemaker_session,\n",
    "    #role=\"arn:aws:iam::<>:role/Admin\"\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "98132b4e-66ed-4d9b-ab5c-a62deeede96b",
   "metadata": {},
   "source": [
    "### Discover and update Finetuning options\n",
    "\n",
    "Each of the technique and model has overridable hyperparameters that can be finetuned by the user."
   ]
  },
  {
   "cell_type": "code",
   "id": "de183042-bb92-4947-9acd-78d7231bda13",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "print(\"Default Finetuning Options:\")\n",
    "pprint(sft_trainer.hyperparameters.to_dict()) # rename as hyperparameters"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "6b57838f-81ac-4fbe-9ddf-5588e42bcce1",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "# To update any hyperparameter, simply assign the value, example:\n",
    "sft_trainer.hyperparameters.global_batch_size=16"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "7453f6a3-e587-4469-8c69-9c207f72b7c8",
   "metadata": {},
   "source": [
    "#### Start SFT training\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "4d3b6441-9abb-447b-9307-9606a8c0fabd",
   "metadata": {
    "scrolled": true,
    "jupyter": {
     "is_executing": true
    }
   },
   "source": [
    "training_job = sft_trainer.train(\n",
    "    wait=True,\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "6e9060b2-e269-461c-88fb-dac4e8854b8a",
   "metadata": {},
   "source": [
    "training_job = sft_trainer.train(\n",
    "    wait=True,\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "0373cea6-7419-47f1-a59e-1cb441324dc3",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "import json\n",
    "import re\n",
    "from sagemaker.core.utils.utils import Unassigned\n",
    "from sagemaker.core.resources import TrainingJob\n",
    "\n",
    "response = TrainingJob.get(training_job_name=\"meta-textgeneration-llama-3-2-1b-instruct-sft-20251201114921\")\n",
    "\n",
    "def pretty_print(obj):\n",
    "    def parse_unassigned(item):\n",
    "        if isinstance(item, Unassigned):\n",
    "            return None\n",
    "        if isinstance(item, dict):\n",
    "            return {k: parse_unassigned(v) for k, v in item.items() if parse_unassigned(v) is not None}\n",
    "        if isinstance(item, list):\n",
    "            return [parse_unassigned(x) for x in item if parse_unassigned(x) is not None]\n",
    "        if isinstance(item, str) and \"Unassigned object\" in item:\n",
    "            pairs = re.findall(r\"(\\w+)=([^<][^=]*?)(?=\\s+\\w+=|$)\", item)\n",
    "            result = {k: v.strip(\"'\\\"\") for k, v in pairs}\n",
    "            return result if result else None\n",
    "        return item\n",
    "\n",
    "    cleaned = parse_unassigned(obj.__dict__ if hasattr(obj, '__dict__') else obj)\n",
    "    print(json.dumps(cleaned, indent=2, default=str))\n",
    "\n",
    "pretty_print(response)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "e9ee7f8e-b26c-4579-9dbc-f08124f2e944",
   "metadata": {},
   "source": [
    "#In order to skip waiting and monitor the training Job later\n",
    "\n",
    "'''\n",
    "training_job = sft_trainer.train(\n",
    "    wait=False,\n",
    ")\n",
    "'''"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "0d99f212-f0bd-43c1-be21-30202fb4a152",
   "metadata": {
    "scrolled": true
   },
   "source": "pretty_print(training_job)",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "da489ad0-36b8-44e7-9f65-2ffd359e5225",
   "metadata": {},
   "source": [
    "### View any Training job details\n",
    "\n",
    "We can get any training job details and its status with TrainingJob.get(...)"
   ]
  },
  {
   "cell_type": "code",
   "id": "6bbe96b4-c8cd-4de3-b4c0-a66fd3086eb2",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "from sagemaker.core.resources import TrainingJob\n",
    "\n",
    "response = TrainingJob.get(training_job_name=\"meta-textgeneration-llama-3-2-1b-instruct-sft-20251123162832\")\n",
    "pretty_print(response)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "566a9dac-343e-446b-83b6-88540bea1766",
   "metadata": {},
   "source": [
    "## Continued Finetuning (or) Finetuning on Model Artifacts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb1ca47e-e373-4436-a568-b4285526edc2",
   "metadata": {},
   "source": [
    "#### Discover a ModelPackage and get its details"
   ]
  },
  {
   "cell_type": "code",
   "id": "11a16b70-526f-42a1-8d0f-3a8b14f559a5",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "from rich import print as rprint\n",
    "from rich.pretty import pprint\n",
    "from sagemaker.core.resources import ModelPackage, ModelPackageGroup\n",
    "\n",
    "#model_package_iter = ModelPackage.get_all(model_package_group_name=\"test-finetuned-models-gamma\")\n",
    "model_package = ModelPackage.get(model_package_name=\"arn:aws:sagemaker:us-west-2:<>:model-package/sdk-test-finetuned-models/2\")\n",
    "\n",
    "pretty_print(model_package)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "ab9705fd-5899-40a7-b4b3-91d3ef1718b8",
   "metadata": {},
   "source": [
    "#### Create Trainer\n",
    "\n",
    "Trainer creation is same as above Finetuning Section except for `model`'s input is ModelPackage(previously trained artifacts)"
   ]
  },
  {
   "cell_type": "code",
   "id": "dc715f3d-543a-4c75-888e-24220d226526",
   "metadata": {},
   "source": [
    "# For fine-tuning \n",
    "sft_trainer = SFTTrainer(\n",
    "    model=model_package, # Union[str, ModelPackage]\n",
    "    training_type=TrainingType.LORA, \n",
    "    model_package_group_name=\"sdk-test-finetuned-models\", # Make it Optional\n",
    "    #mlflow_resource_arn=\"arn:aws:sagemaker:us-west-2:<>:mlflow-tracking-server/mmlu-eval-experiment\",  # Optional[str] - MLflow app ARN (auto-resolved if not provided), can accept name and search in the account\n",
    "    mlflow_experiment_name=\"test-finetuned-models-exp\", # Optional[str]\n",
    "    mlflow_run_name=\"test-finetuned-models-run\", # Optional[str]\n",
    "    training_dataset=\"s3://mc-flows-sdk-testing/input_data/sft/\", #Optional[]\n",
    "    s3_output_path=\"s3://mc-flows-sdk-testing/output/\",\n",
    ")\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "293a462a-ae17-47d0-80a1-d0c5b54a9cfc",
   "metadata": {},
   "source": [
    "#### Start the Training"
   ]
  },
  {
   "cell_type": "code",
   "id": "5c3f3380-5305-4178-9120-aeca1ba6ea44",
   "metadata": {},
   "source": [
    "training_job = sft_trainer.train(\n",
    "    wait=True,\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "d55299f5-639b-435c-8b50-718491c0060a",
   "metadata": {},
   "source": "pretty_print(training_job)",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "00ae0223-2e35-4a69-aa27-c96665b35172",
   "metadata": {},
   "source": [
    "#### SFT Trainer Nova testing"
   ]
  },
  {
   "cell_type": "code",
   "id": "0f11e6d5-7bb9-41b8-8d6c-94377691e3be",
   "metadata": {},
   "source": [
    "os.environ['SAGEMAKER_REGION'] = 'us-east-1'\n",
    "\n",
    "# For fine-tuning \n",
    "sft_trainer_nova = SFTTrainer(\n",
    "    #model=\"test-nova-lite-v2\", # Union[str, ModelPackage]\n",
    "    #model=\"nova-textgeneration-micro\",\n",
    "    model=\"nova-textgeneration-lite-v2\",\n",
    "    training_type=TrainingType.LORA, \n",
    "    model_package_group_name=\"sdk-test-finetuned-models\", # Make it Optional\n",
    "    #mlflow_resource_arn=\"arn:aws:sagemaker:us-east-1:<>:mlflow-app/app-UNBKLOAX64PX\",  # Optional[str] - MLflow app ARN (auto-resolved if not provided), can accept name and search in the account\n",
    "    mlflow_experiment_name=\"test-nova-finetuned-models-exp\", # Optional[str]\n",
    "    mlflow_run_name=\"test-nova-finetuned-models-run\", # Optional[str]\n",
    "    training_dataset=\"arn:aws:sagemaker:us-east-1:<>:hub-content/sdktest/DataSet/sft-nova-test-dataset/0.0.1\",\n",
    "    s3_output_path=\"s3://mc-flows-sdk-testing-us-east-1/output/\"\n",
    ")\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "9acf53aa-4ac5-4a74-8d67-607e0d09820f",
   "metadata": {},
   "source": [
    "sft_trainer_nova.hyperparameters.to_dict()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "e6ce0ece-0185-4fc4-ae89-218667cf6b14",
   "metadata": {},
   "source": [
    "training_job = sft_trainer_nova.train(\n",
    "    wait=True,\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
